from collections import namedtuple

TRANSCRIPT = namedtuple('TRANSCRIPT', ['name', 'chrom', 'tx_start', 'tx_end', 'strand', 'exons', 'introns'])


def read_refgene(infile):
    """
    Iterate through a refGene file.

    GenePred extension format:
    http://genome.ucsc.edu/FAQ/FAQformat.html#GenePredExt

    Column definitions:
    0. uint undocumented id
    1. string name;                 "Name of gene (usually transcript_id from GTF)"
    2. string chrom;                "Chromosome name"
    3. char[1] strand;              "+ or - for strand"
    4. uint txStart;                "Transcription start position"
    5. uint txEnd;                  "Transcription end position"
    6. uint cdsStart;               "Coding region start"
    7. uint cdsEnd;                 "Coding region end"
    8. uint exonCount;              "Number of exons"
    9. uint[exonCount] exonStarts;  "Exon start positions"
    10. uint[exonCount] exonEnds;   "Exon end positions"
    11. int score;                  "Score"
    12. string name2;               "Alternate name (e.g. gene_id from GTF)"
    13. string cdsStartStat;        "enum('none','unk','incmpl','cmpl')"
    14. string cdsEndStat;          "enum('none','unk','incmpl','cmpl')"
    15. lstring exonFrames;         "Exon frame offsets {0,1,2}"
    """
    with open(infile) as fi:
        for line in fi:
            if line.startswith('#'):
                continue
            row = line.rstrip('\n').split('\t')
            exon_starts = map(int, row[9].strip(',').split(','))
            exon_ends = map(int, row[10].strip(',').split(','))
            exons = list(zip(exon_starts, exon_ends))

            yield {
                'bin': row[0],
                'name': row[1],
                'chrom': row[2],
                'strand': row[3],
                'txStart': int(row[4]),
                'txEnd': int(row[5]),
                'cdsStart': int(row[6]),
                'cdsEnd': int(row[7]),
                'exonCount': int(row[8]),
                'exonStarts': [int(i) for i in row[9].strip(',').split(',')],
                'exonEnds': [int(i) for i in row[10].strip(',').split(',')],
                'score': int(row[11]),
                'name2': row[12],
                'cdsStartStat': row[13],
                'cdsEndStat': row[14],
                'exonFrames': row[15],
                'exons': exons,
            }


def read_transcripts(infile) -> dict:
    trans_dict = dict()
    for trans in read_refgene(infile):
        exons = trans['exons']
        introns = [(exons[i][1], exons[i+1][0]) for i in range(len(exons) - 1)]
        transcript = TRANSCRIPT(
            name=trans['name'], chrom=trans['chrom'], strand=trans['strand'],
            tx_start=trans['txStart'], tx_end=trans['txEnd'], exons=exons, introns=introns
        )
        trans_dict[transcript.name] = transcript
    return trans_dict
